package com.example.pt.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.*;

/**
 * BT种子Bencode格式处理工具类
 */
public class BencodingUtils {
    private static final Logger logger = LoggerFactory.getLogger(BencodingUtils.class);

    /**
     * 读取并解析BT种子文件
     * @param file 种子文件
     * @return 解析后的字典对象
     */
    public static Map<String, Object> decodeTorrentFile(File file) throws IOException {
        try (BufferedInputStream bis = new BufferedInputStream(new FileInputStream(file))) {
            return (Map<String, Object>) decodeNextObject(bis);
        }
    }
    
    /**
     * 将种子对象保存为文件
     * @param torrentData 种子数据
     * @param outputFile 输出文件
     */
    public static void encodeTorrentFile(Map<String, Object> torrentData, File outputFile) throws IOException {
        try (OutputStream os = new FileOutputStream(outputFile)) {
            encodeObject(torrentData, os);
        }
    }

    /**
     * 获取种子文件中描述的实际内容总大小
     * @param torrentFile 种子文件
     * @return 内容总大小（字节）
     */
    public static long getTotalContentSize(File torrentFile) {
        try {
            // 解析种子文件
            Map<String, Object> torrentData = decodeTorrentFile(torrentFile);
            
            // 获取info字典
            Map<String, Object> info = (Map<String, Object>) torrentData.get("info");
            if (info == null) {
                logger.error("无法获取种子的info字典");
                return 0;
            }
            
            long totalSize = 0;
            
            // 检查是否是多文件种子
            if (info.containsKey("files")) {
                // 多文件模式
                List<Map<String, Object>> files = (List<Map<String, Object>>) info.get("files");
                for (Map<String, Object> fileDict : files) {
                    long length = (Long) fileDict.get("length");
                    totalSize += length;
                }
            } else if (info.containsKey("length")) {
                // 单文件模式
                totalSize = (Long) info.get("length");
            }
            
            return totalSize;
        } catch (Exception e) {
            logger.error("获取种子内容大小失败", e);
            return 0;
        }
    }

    /**
     * 更新种子文件中的Tracker URL，添加或替换passkey
     * @param torrentFile 种子文件
     * @param outputFile 输出文件
     * @param passkey 用户passkey
     * @return 是否成功
     */
    public static boolean updatePasskey(File torrentFile, File outputFile, String passkey) {
        try {
            // 解析种子文件
            Map<String, Object> torrentData = decodeTorrentFile(torrentFile);
            
            // 获取announce URL
            String announceUrl = new String((byte[]) torrentData.get("announce"), StandardCharsets.ISO_8859_1);
            logger.info("原始announce URL: {}", announceUrl);
            
            // 更新announce URL
            String newAnnounceUrl = updateUrlWithPasskey(announceUrl, passkey);
            logger.info("新的announce URL: {}", newAnnounceUrl);
            
            // 替换announce URL
            torrentData.put("announce", newAnnounceUrl.getBytes(StandardCharsets.ISO_8859_1));
            
            // 如果存在announce-list，也更新它
            if (torrentData.containsKey("announce-list")) {
                List<Object> announceList = (List<Object>) torrentData.get("announce-list");
                
                for (int i = 0; i < announceList.size(); i++) {
                    List<Object> tier = (List<Object>) announceList.get(i);
                    
                    for (int j = 0; j < tier.size(); j++) {
                        String url = new String((byte[]) tier.get(j), StandardCharsets.ISO_8859_1);
                        String newUrl = updateUrlWithPasskey(url, passkey);
                        tier.set(j, newUrl.getBytes(StandardCharsets.ISO_8859_1));
                    }
                }
            }
            
            // 保存修改后的种子文件
            encodeTorrentFile(torrentData, outputFile);
            
            return true;
        } catch (Exception e) {
            logger.error("更新种子文件passkey失败", e);
            return false;
        }
    }
    
    /**
     * 设置种子文件的Tracker URL，可以清空或设置为指定URL
     * @param torrentFile 输入种子文件
     * @param outputFile 输出种子文件
     * @param trackerUrl 要设置的tracker URL，如果为null则清空tracker
     * @param passkey 可选的passkey，如果提供则添加到tracker URL
     * @return 是否成功
     */
    public static boolean setTrackerUrl(File torrentFile, File outputFile, String trackerUrl, String passkey) {
        try {
            // 解析种子文件
            Map<String, Object> torrentData = decodeTorrentFile(torrentFile);
            
            if (trackerUrl == null || trackerUrl.isEmpty()) {
                // 清空tracker URL
                logger.info("清空种子文件的tracker URL");
                torrentData.put("announce", "".getBytes(StandardCharsets.ISO_8859_1));
                
                // 移除announce-list
                if (torrentData.containsKey("announce-list")) {
                    torrentData.remove("announce-list");
                    logger.info("已移除announce-list");
                }
            } else {
                // 设置新的tracker URL
                String finalUrl = trackerUrl;
                
                // 如果提供了passkey，添加到URL
                if (passkey != null && !passkey.isEmpty()) {
                    if (trackerUrl.contains("?")) {
                        finalUrl = trackerUrl + "&passkey=" + passkey;
                    } else {
                        finalUrl = trackerUrl + "?passkey=" + passkey;
                    }
                }
                
                logger.info("设置新的tracker URL: {}", finalUrl);
                torrentData.put("announce", finalUrl.getBytes(StandardCharsets.ISO_8859_1));
                
                // 清除announce-list，仅使用单一tracker
                if (torrentData.containsKey("announce-list")) {
                    torrentData.remove("announce-list");
                    logger.info("已移除announce-list，仅使用单一tracker");
                }
            }
            
            // 保存修改后的种子文件
            encodeTorrentFile(torrentData, outputFile);
            
            return true;
        } catch (Exception e) {
            logger.error("设置种子文件tracker URL失败", e);
            return false;
        }
    }
    
    /**
     * 解析下一个对象
     */
    private static Object decodeNextObject(InputStream is) throws IOException {
        is.mark(1);
        int ch = is.read();
        if (ch < 0) {
            throw new EOFException("Unexpected end of stream");
        }
        
        switch (ch) {
            case 'i': // 整数
                return decodeInteger(is);
            case 'l': // 列表
                return decodeList(is);
            case 'd': // 字典
                return decodeDictionary(is);
            case '0': case '1': case '2': case '3': case '4':
            case '5': case '6': case '7': case '8': case '9':
                // 字符串（带长度前缀）
                is.reset();
                return decodeString(is);
            default:
                throw new IOException("Invalid bencode format: " + (char)ch);
        }
    }
    
    /**
     * 解析整数
     */
    private static Long decodeInteger(InputStream is) throws IOException {
        StringBuilder sb = new StringBuilder();
        int ch;
        while ((ch = is.read()) >= 0 && ch != 'e') {
            sb.append((char)ch);
        }
        return Long.parseLong(sb.toString());
    }
    
    /**
     * 解析字符串
     */
    private static byte[] decodeString(InputStream is) throws IOException {
        StringBuilder sb = new StringBuilder();
        int ch;
        // 读取长度
        while ((ch = is.read()) >= 0 && ch != ':') {
            sb.append((char)ch);
        }
        
        int length = Integer.parseInt(sb.toString());
        byte[] result = new byte[length];
        int read = 0;
        while (read < length) {
            int count = is.read(result, read, length - read);
            if (count < 0) {
                throw new EOFException("Unexpected end of stream");
            }
            read += count;
        }
        return result;
    }
    
    /**
     * 解析列表
     */
    private static List<Object> decodeList(InputStream is) throws IOException {
        List<Object> list = new ArrayList<>();
        while (true) {
            is.mark(1);
            int ch = is.read();
            if (ch == 'e') {
                break;
            }
            is.reset();
            list.add(decodeNextObject(is));
        }
        return list;
    }
    
    /**
     * 解析字典
     */
    private static Map<String, Object> decodeDictionary(InputStream is) throws IOException {
        Map<String, Object> dict = new HashMap<>();
        while (true) {
            is.mark(1);
            int ch = is.read();
            if (ch == 'e') {
                break;
            }
            is.reset();
            
            // 键必须是字符串
            byte[] keyBytes = decodeString(is);
            String key = new String(keyBytes, StandardCharsets.ISO_8859_1);
            
            // 解析值
            Object value = decodeNextObject(is);
            dict.put(key, value);
        }
        return dict;
    }
    
    /**
     * 编码对象到输出流
     */
    private static void encodeObject(Object obj, OutputStream os) throws IOException {
        if (obj instanceof Long || obj instanceof Integer) {
            encodeInteger((Number)obj, os);
        } else if (obj instanceof byte[]) {
            encodeString((byte[])obj, os);
        } else if (obj instanceof String) {
            encodeString(((String)obj).getBytes(StandardCharsets.ISO_8859_1), os);
        } else if (obj instanceof List) {
            encodeList((List<Object>)obj, os);
        } else if (obj instanceof Map) {
            encodeDictionary((Map<String, Object>)obj, os);
        } else {
            throw new IOException("Unsupported object type: " + obj.getClass().getName());
        }
    }
    
    /**
     * 编码整数
     */
    private static void encodeInteger(Number n, OutputStream os) throws IOException {
        os.write('i');
        os.write(n.toString().getBytes(StandardCharsets.ISO_8859_1));
        os.write('e');
    }
    
    /**
     * 编码字符串
     */
    private static void encodeString(byte[] s, OutputStream os) throws IOException {
        os.write(String.valueOf(s.length).getBytes(StandardCharsets.ISO_8859_1));
        os.write(':');
        os.write(s);
    }
    
    /**
     * 编码列表
     */
    private static void encodeList(List<Object> list, OutputStream os) throws IOException {
        os.write('l');
        for (Object item : list) {
            encodeObject(item, os);
        }
        os.write('e');
    }
    
    /**
     * 编码字典
     */
    private static void encodeDictionary(Map<String, Object> dict, OutputStream os) throws IOException {
        os.write('d');
        
        // 按键名升序排序
        TreeMap<String, Object> sortedDict = new TreeMap<>(dict);
        
        for (Map.Entry<String, Object> entry : sortedDict.entrySet()) {
            byte[] keyBytes = entry.getKey().getBytes(StandardCharsets.ISO_8859_1);
            encodeString(keyBytes, os);
            encodeObject(entry.getValue(), os);
        }
        os.write('e');
    }

    /**
     * 更新URL添加或替换passkey
     */
    private static String updateUrlWithPasskey(String url, String passkey) {
        if (url.contains("passkey=")) {
            // 替换现有passkey
            return url.replaceAll("passkey=[^&]+", "passkey=" + passkey);
        } else if (url.contains("?")) {
            // 有查询参数，添加passkey
            return url + "&passkey=" + passkey;
        } else {
            // 没有查询参数，添加passkey
            return url + "?passkey=" + passkey;
        }
    }

    /**
     * 获取种子文件的InfoHash值
     * @param torrentFile 种子文件
     * @return InfoHash值（十六进制字符串）
     */
    public static String getInfoHash(File torrentFile) {
        try {
            // 解析种子文件
            Map<String, Object> torrentData = decodeTorrentFile(torrentFile);
            
            // 获取info字典
            Map<String, Object> info = (Map<String, Object>) torrentData.get("info");
            if (info == null) {
                logger.error("无法获取种子的info字典");
                return null;
            }
            
            // 重新编码info字典以计算其SHA1哈希值
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            encodeObject(info, baos);
            byte[] infoBytes = baos.toByteArray();
            
            // 计算SHA1哈希
            MessageDigest sha1 = MessageDigest.getInstance("SHA-1");
            byte[] hash = sha1.digest(infoBytes);
            
            // 转换为十六进制字符串
            StringBuilder hexString = new StringBuilder();
            for (byte b : hash) {
                String hex = Integer.toHexString(0xff & b);
                if (hex.length() == 1) {
                    hexString.append('0');
                }
                hexString.append(hex);
            }
            
            return hexString.toString();
        } catch (Exception e) {
            logger.error("获取种子InfoHash失败", e);
            return null;
        }
    }
} 